MUM 2024-25 Modele dyfuzyjne

%matplotlib inline
#Required for Jupyter to display plots as output 
# inline -→> display inLine for Jupyter
# notebook -=› display inLine, enabling interactivity. (Awesome.) import matplotlib as mpl
import matplotlib.pyplot as plt

Lil'Log Weng, Lilian. (Jul 2021). What are diffusion models? Lil’Log. https://lilianweng.github.io/posts/2021-07-11-diffusion-models/. SuperAnnotate: introduction to diffusion models for ML

modele dyfuzyjne

Wszystkie modele mają problemy, np. GANy są niestabilne w uczeniu.

Modele dyfuzyjne są inspirowane modelami termodynamiki, które nie są w stanie stabilnym (non-equilibrium thermodynamics).

Definiują łańcuch Markowa wiele kroków dyfuzji przez powolne dodawanie szumu do danych, a poźniej uczą się procesu odwrotnego do rekonstrukcji.

[Fig. Encord]

Modele są uczone według ustalonej procedury, nie stochastycznie, z przestrzenią ukrytą o wymiarze danych. Proces jest dwu fazowy: w przód (forward) i z porotem (reverse).

Proces forward

Forward

• Dla $x_0\sim\,q(x)$, proces dyfuzji forward tworzy zaszumione próbki $x_1,\dots,x_T$
• Kroki są kontrolowane schematem wariancji $\{\beta_t\in(0,1)\}_{t=1}^T$

$$q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t}\, \mathbf{x}_{t-1}, \beta_t\mathbf{I}) \quad q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod^T_{t=1} q(\mathbf{x}_t \vert \mathbf{x}_{t-1}),$$ gdzie $\mu_t=\sqrt{1 - \beta_t} \mathbf{x}_{t-1}$ staje się średnią a $\beta_t$ rośnie.

• Próbki coraz bardziej tracą wyraźne cechy
• Proces odwrotny $q(x_{t-1}\mid x_t)$ jest nieznany $$\begin{aligned} \mathbf{x}_t &= \sqrt{1-\beta_t}\mathbf{x}_{t-1} + \sqrt{1 - (1-\beta_t)}\boldsymbol{\epsilon}_{t-1} \end{aligned} $$ niech $\alpha_t=1-\beta_t, \overline\alpha_t=\prod_{t=1}^t\alpha_t, \boldsymbol{\epsilon}_{t-1}, \boldsymbol{\epsilon}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ $$ \begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\boldsymbol{\epsilon}_{t-1} & \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\boldsymbol{\epsilon}}_{t-2} & \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon} \end{aligned} $$ stąd całą zależność $x_t$ od $x_0$ można zapisać $$ \begin{aligned} q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \end{aligned} $$

kolejne kroki

• Łącząc dwa rozkłady normalne z różną wariancją, dostajemy odchylenie $$\mathcal{N}(0,(\sigma_1^2+\sigma_2^2)\mathbb{I})$$ a stąd tutaj $$\sqrt{(1-\alpha_t)+\alpha_t(1-\alpha_{t-1})}=\sqrt{1-\alpha_t\alpha_{t-1}}$$ a po wielu krokach odchylenie $\sqrt{1-\prod_{t=1}^T}=\sqrt{1-\bar{\alpha}_T}$.
• to przypomina modyfikację nauczania SGD gdzie do gradientu dodawany jest jeszcze szum utrudniający utknięcie procesu w płytkim lokalnym minimum.
• w procesie wykorzystywany jest ciąg $\beta_t=1-\alpha_t$,
  • zwykle od $\beta_1=10^{-4}$ do $\beta_T=0.02$.
• jest szereg modyfikacji, które poprawiają znacznie zbieżność. Szereg wariancji może być ustalony, może jednak być uczony w trakcie (pod warunkiem wzrostu).

Proces odwrotny (reverse)

• Dla bardzo dużych $T$ rozkład coraz bardziej będzie przypominać $\mathcal{N}(0, \mathbb{I})$,
  • stąd rozpoczynając proces odwrotny będzie można z tego rozkładu losować

Gausowskość

• Dla wystarczająco małych $\beta_t$, wartość $q(x_{t-1}\mid x_t)$ także będzie Gausowska.

• jednak ta estymacja wymagałaby znajomości całego zbioru danych.

• Możemy jednak estymować $$p_\theta(x_{0:T})=p(x_T)\prod_{t=1}^Tp_\theta(x_{t-1}\mid x_t)$$ dla $$p_\theta(x_{t-1}\mid x_t)=\mathcal{N}(x_{t-1}; \mu_\theta(x_t,t),\Sigma_\theta(x_t,t))$$

• uczenie polega wiec na znalezieniu takiego $p_\theta()$, który będzie aproksymował proces odwrotny $$ \begin{aligned} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\ \end{aligned} $$ i wymagać będzie wyznaczenia odpowiednich $\mu, \epsilon_t$.

• Okazuje się, że proces ma uwarunkowania bardzo podobne do VAE,

  • stad można użyć dolnego ograniczenia ELBO i odpowiednio przeprowadzić nauczanie.

Koszt kolejnych kroków i optymalizacja

• Koszt można przedtawić jako sumę kosztów poszczególnych kroków $$L_\text{VLB} = L_T + L_{T-1} + \dots + L_0,$$ gdzie $$ \begin{aligned} L_T &= KL(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T)) \\ L_t &= KL(q(\mathbf{x}_t \vert \mathbf{x}_{t+1}, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_t \vert\mathbf{x}_{t+1})) \text{ for }1 \leq t \leq T-1 \\ L_0 &= - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1). \end{aligned} $$
• Wszystkie elementy dywergencji są przedstawione w postaci zamkniętej,
  • nie wymagają samplowania dla porównania rozkładów.
• Dla $L_0$ potrzebny jest dodatkowy dekoder.
• Poszczególne elementy $L_t$ są możliwe do sparametryzowania przez $\mu_t$ oraz $\Sigma_\theta$ i ich predykcję.
• Okazuje się, ze nawet uproszczony model tej predykcji działa lepiej i stabilniej.

Inne procesy odszumiania

• Ten proces jest podobny do procesów odszumiania, np. Denoising Auto Encoder DAE.
• W DAE do przykładów jest dodawany szum, by wymusić jego redukcję, a przez to nauczenie się prawdziwego rozkładu danych.

Warunkowy proces generowania

• uczenie klasyfikatora $f_\theta(y\mid x_t,t)$ na zaszumionym obrazie

• wykorzystuje gradient $$\nabla_x\log\,f_\theta(y\mid x_t)$$ sterujący procesem samplowania w kierunku etykiety warunkującej $y$.

• Proces jest wzmocnieniem samego samplingu.

  • Algorytm ablated diffusion model ADM osiągnął lepsze wyniki niż najlepszy duży model BigGAN.

• Bez explicite nauczonego klasyfikatora też jest możliwe uzyskanie warunkowej generacji. Możemy mieć

1. bezwarunkowy proces odszumiający $p\theta(x)$ zbudowany przez estymator zysku (score) $\epsilon_\theta(x_t,t)$ oraz
2. model warunkowy $p\theta(x\mid y)$.

• Te modele mogą być nauczone przez jedną sieć neuronową tak, że $p_\theta(x\mid y)$ jest uczony na parach $(x, y)$.
• Jednak wartość warunkowa $y$ jest losowo usuwana tak, ze model będzie w stanie generować obiekty także bezwarunkowo.
• Taki model osiąga dobry balans między FID (dobre cechy i rozróżnienie miedzy obiektami prawdziwymi i generowanymi) oraz IS (zróżnicowanie i jakość).
• Model GLIDE wykorzystywał model CLIP jako klasyfikator oraz sterowanie bez klasyfikatora znajdując drugie podejście jako zwykle lepsze.

Przyspieszenie procesu

• Cały proces tworzenia modelu może być bardzo wolny, gdyż liczba kroków $T$ może być bardzo duża.
• Można zmniejszyć liczbę kroków, poprawiając model co $n$ oryginalnych kroków.

50 tysięcy obrazów 32x32 może zająć 20 godzin dla procesu Markowa, ale mniej niż minutę dla modelu GAN.

Destylacja

• W procesie uczenia modelu (w przód, forward) można wykorzystać destylację
  • proces uczy się kolejnych kroków $q(x_t\mid x_{t-1}$
  • po nauczeniu dwóch kolejnych następuje destylacja
    • dwa kolejne służą jako nauczyciel
    • uczą jeden kolejny proces, który bedzie przyspieszał pracę

Skalowanie rozdzielczości oraz jakości

• Możliwością jest stworzenie strumienia (pipeline) przez utworzenie szeregu modeli dyfuzji w coraz wyższych rozdzielczościach.
• Podstawowym elementem jest augmentacja oparta szumem pomiędzy elementami.
• Będzie to polegać na dodaniu silnej augmentacji danych na warunkującym wejściu $z$ dla każdego $p_\theta(x\mid z)$.
• bardzo częstym wyborem jest U-net (o tym poniżej). ![[cascaded-diffusion.png]] [Rysunek Ho] Najlepiej dodać szum Gausowski dla niskich rozdzielczości , oraz Gaussian blur w wysokich (zamiast uśrednionego, daje większą wagę centralnym pikselom; inaczej Gausowskie wygładzanie).

1. Dla obrazu niskiej rozdzielczości proces jest zatrzymywany dla niskiego $t$.
2. Odwrotny proces niskiej rozdzielczości dochodzi do punktu $0$
3. Wtedy jest zaburzany przez $z_t\sim q(x_t\mid x_0)$, a ten podaje $z_t$ do modelu wyższej rozdzielczości.

Model unCLIP

• Przy generowaniu obrazów sterowanych tekstem wydajne jest użycie CLIP.

• Dla CLIP uczonego parami $(x, y)$, gdzie $x$ to obraz, a $y$ to związana etykieta,

  • budujemy embeddingi tekstu $c^t(y)$ i obrazu $c^i(x)$. unCLIP uczy się równolegle modeli
  • model prior $P(c^i\mid y)$ dla zadanego $y$ zwraca embedding (dla CLIP) $c^i$ tekstu $y$,
  • dekoder $P(x\mid x^i, [y])$ generuje obraz $x$ mając embedding $c^i$; ewentualnie oryginalny tekst $y$.

• Modele pozwalają na generowanie warunkowe $$P(x\mid y) = P(x, c^i\mid y) = P(x\mid c^i,y)\,P(c^i\mid y).$$ [Ramesh 2022]

• Niektóre inne modele wykorzystują LLM dla kodowaniaopisu obrazu przy jego generowaniu.

  • Większa waga dla dobrego dopasowania tekstu prowadzi często do słabszej jakości obrazu. Konieczny jest często tuning.

Architektura modeli

Podstawowymi dla modeli dyfuzyjnych są U-Net i Transformer.

• U-Net [Ronneberger]
• Transformer

Transformer

Tu tworzony jest Diffusion Transformer (DiT), który

• jako wejście bierze reprezentację typu latent,
• ta jest dzielona na $p^2$ kawałków (patchy), jak w ViT,
• sekwencja tych patchy jest podawana na wejście transformera,
• między embeddingiem a wartwą normalizacji warstwowej (LayerNorm) i wartwą liniową znajduje się dodatkowy blok DiT,
  • blok składa przypomina blok Transformera i przewiduje parametry potrzebne do wygenerowania szumu kolejnej transformacji dyfuzji,
• bloków Transformera, z dodanym DiT, jest N. Model Transformera jest wyjątkowo elastyczny, jednak w połączeniu z koniecznym długim łańcuchem Markowa w przetwarzaniu dyfuzji, może być kosztowny.

Przykładowe uczenie (denoising-diffusion-pytorch)

import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion

/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/io/image.py:14: UserWarning: Failed to load image Python extension: 'dlopen(/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/image.so, 0x0006): Library not loaded: @rpath/libjpeg.9.dylib
  Referenced from: <0B7EB158-53DC-3403-8A49-22178CAB4612> /Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/image.so
  Reason: tried: '/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/../../../libjpeg.9.dylib' (no such file), '/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/../../../libjpeg.9.dylib' (no such file), '/Users/igor/miniforge3/envs/mini/lib/python3.10/lib-dynload/../../libjpeg.9.dylib' (no such file), '/Users/igor/miniforge3/envs/mini/bin/../lib/libjpeg.9.dylib' (no such file)'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(

# Definicja modelu Unet
# info wymiar map cech oraz mnożniki dla kolejnych kompresji
model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)

# Model dyfuzji
diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
#    loss_type = 'l1'    # L1 or L2
)

# uczenie z szumu
training_images = torch.randn(8, 3, 128, 128)
loss = diffusion(training_images)
loss.backward()

# generowanie obrazów

sampled_images = diffusion.sample(batch_size = 4)

sampling loop time step:   0%|          | 0/1000 [00:00<?, ?it/s]

import torchvision

to_pil = torchvision.transforms.ToPILImage()

def show(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')

for k in range(4):
    show(sampled_images[k])

# uczenie dla zbioru danych
from denoising_diffusion_pytorch import Unet, GaussianDiffusion
from denoising_diffusion_pytorch import Trainer

model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8)
)

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,   # number of steps
    # loss_type = 'l1'    # L1 or L2
)

trainer = Trainer(
    diffusion,
    '/Users/igor/tmp',
    train_batch_size = 32,
    train_lr = 2e-5,
    train_num_steps = 700000,         # total training steps
    gradient_accumulate_every = 2,    # gradient accumulation steps
    ema_decay = 0.995,                # exponential moving average decay
    amp = True                        # turn on mixed precision
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[38], line 1
----> 1 trainer = Trainer(
      2     diffusion,
      3     '/Users/igor/tmp',
      4     train_batch_size = 32,
      5     train_lr = 2e-5,
      6     train_num_steps = 700000,         # total training steps
      7     gradient_accumulate_every = 2,    # gradient accumulation steps
      8     ema_decay = 0.995,                # exponential moving average decay
      9     amp = True                        # turn on mixed precision
     10 )

File ~/miniforge3/envs/mini/lib/python3.10/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py:910, in Trainer.__init__(self, diffusion_model, folder, train_batch_size, gradient_accumulate_every, augment_horizontal_flip, train_lr, train_num_steps, ema_update_every, ema_decay, adam_betas, save_and_sample_every, num_samples, results_folder, amp, mixed_precision_type, split_batches, convert_image_to, calculate_fid, inception_block_idx, max_grad_norm, num_fid_samples, save_best_and_latest_only)
    906 super().__init__()
    908 # accelerator
--> 910 self.accelerator = Accelerator(
    911     split_batches = split_batches,
    912     mixed_precision = mixed_precision_type if amp else 'no'
    913 )
    915 # model
    917 self.model = diffusion_model

File ~/miniforge3/envs/mini/lib/python3.10/site-packages/accelerate/accelerator.py:485, in Accelerator.__init__(self, device_placement, split_batches, mixed_precision, gradient_accumulation_steps, cpu, dataloader_config, deepspeed_plugin, fsdp_plugin, megatron_lm_plugin, rng_types, log_with, project_dir, project_config, gradient_accumulation_plugin, step_scheduler_with_optimizer, kwargs_handlers, dynamo_backend, deepspeed_plugins)
    481 self.native_amp = True
    482 if self.device.type not in ("xpu", "cuda", "npu", "xla", "mlu", "musa") or is_torch_xla_available(
    483     check_is_tpu=True
    484 ):
--> 485     raise ValueError(f"fp16 mixed precision requires a GPU (not {self.device.type!r}).")
    486 kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {}
    487 self.scaler = get_grad_scaler(self.distributed_type, **kwargs)

ValueError: fp16 mixed precision requires a GPU (not 'mps').

trainer.train()

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[39], line 1
----> 1 trainer.train()

NameError: name 'trainer' is not defined